package cmd

import (
	"encoding/hex"
	"encoding/json"
	"fmt"
	"path/filepath"
	"strings"
	"time"

	"github.com/onflow/cadence"

	"github.com/spf13/cobra"

	"github.com/onflow/flow-go/cmd"
	"github.com/onflow/flow-go/cmd/bootstrap/run"
	"github.com/onflow/flow-go/cmd/bootstrap/utils"
	"github.com/onflow/flow-go/cmd/util/cmd/common"
	hotstuff "github.com/onflow/flow-go/consensus/hotstuff/model"
	"github.com/onflow/flow-go/fvm"
	model "github.com/onflow/flow-go/model/bootstrap"
	"github.com/onflow/flow-go/model/dkg"
	"github.com/onflow/flow-go/model/flow"
	"github.com/onflow/flow-go/module/epochs"
	"github.com/onflow/flow-go/state/protocol/badger"
	"github.com/onflow/flow-go/state/protocol/inmem"
	"github.com/onflow/flow-go/state/protocol/protocol_state"
	"github.com/onflow/flow-go/state/protocol/protocol_state/kvstore"
	"github.com/onflow/flow-go/utils/io"
)

var (
	flagConfig                  string
	flagInternalNodePrivInfoDir string
	flagPartnerNodeInfoDir      string
	// Deprecated: use flagPartnerWeights instead
	deprecatedFlagPartnerStakes           string
	flagPartnerWeights                    string
	flagDKGDataPath                       string
	flagRootBlockPath                     string
	flagRootCommit                        string
	flagIntermediaryBootstrappingDataPath string
	flagRootBlockVotesDir                 string
	// optional flags for creating
	flagServiceAccountPublicKeyJSON string
	flagGenesisTokenSupply          string
)

// finalizeCmd represents the finalize command`
var finalizeCmd = &cobra.Command{
	Use:   "finalize",
	Short: "Finalize the bootstrapping process",
	Long: `Finalize the bootstrapping process, which includes running the DKG for the generation of the random beacon
	keys and generating the root block, QC, execution result and block seal.`,
	Run: finalize,
}

func init() {
	rootCmd.AddCommand(finalizeCmd)
	addFinalizeCmdFlags()
}

func addFinalizeCmdFlags() {
	// required parameters for network configuration and generation of root node identities
	finalizeCmd.Flags().StringVar(&flagConfig, "config", "",
		"path to a JSON file containing multiple node configurations (fields Role, Address, Weight)")
	finalizeCmd.Flags().StringVar(&flagInternalNodePrivInfoDir, "internal-priv-dir", "", "path to directory "+
		"containing the output from the `keygen` command for internal nodes")
	finalizeCmd.Flags().StringVar(&flagPartnerNodeInfoDir, "partner-dir", "", "path to directory "+
		"containing one JSON file starting with node-info.pub.<NODE_ID>.json for every partner node (fields "+
		" in the JSON file: Role, Address, NodeID, NetworkPubKey, StakingPubKey)")
	// Deprecated: remove this flag
	finalizeCmd.Flags().StringVar(&deprecatedFlagPartnerStakes, "partner-stakes", "", "deprecated: use partner-weights instead")
	finalizeCmd.Flags().StringVar(&flagPartnerWeights, "partner-weights", "", "path to a JSON file containing "+
		"a map from partner node's NodeID to their weight")
	finalizeCmd.Flags().StringVar(&flagDKGDataPath, "dkg-data", "", "path to a JSON file containing data as output from DKG process")
	finalizeCmd.Flags().StringVar(&flagRootCommit, "root-commit", "0000000000000000000000000000000000000000000000000000000000000000", "state commitment of root execution state")

	cmd.MarkFlagRequired(finalizeCmd, "config")
	cmd.MarkFlagRequired(finalizeCmd, "internal-priv-dir")
	cmd.MarkFlagRequired(finalizeCmd, "partner-dir")
	cmd.MarkFlagRequired(finalizeCmd, "partner-weights")
	cmd.MarkFlagRequired(finalizeCmd, "dkg-data")
	cmd.MarkFlagRequired(finalizeCmd, "root-commit")

	// required parameters for generation of root block, root execution result and root block seal
	finalizeCmd.Flags().StringVar(&flagRootBlockPath, "root-block", "", "path to a JSON file containing root block")
	finalizeCmd.Flags().StringVar(&flagIntermediaryBootstrappingDataPath, "intermediary-bootstrapping-data", "", "path to a JSON file containing intermediary bootstrapping data generated by the rootblock command")
	finalizeCmd.Flags().StringVar(&flagRootBlockVotesDir, "root-block-votes-dir", "", "path to directory with votes for root block")

	cmd.MarkFlagRequired(finalizeCmd, "root-block")
	cmd.MarkFlagRequired(finalizeCmd, "intermediary-bootstrapping-data")
	cmd.MarkFlagRequired(finalizeCmd, "root-block-votes-dir")

	// these two flags are only used when setup a network from genesis
	finalizeCmd.Flags().StringVar(&flagServiceAccountPublicKeyJSON, "service-account-public-key-json",
		"{\"PublicKey\":\"ABCDEFGHIJK\",\"SignAlgo\":2,\"HashAlgo\":1,\"SeqNumber\":0,\"Weight\":1000}",
		"encoded json of public key for the service account")
	finalizeCmd.Flags().StringVar(&flagGenesisTokenSupply, "genesis-token-supply", "10000000.00000000",
		"genesis flow token supply")
}

func finalize(cmd *cobra.Command, args []string) {

	// maintain backward compatibility with old flag name
	if deprecatedFlagPartnerStakes != "" {
		log.Warn().Msg("using deprecated flag --partner-stakes (use --partner-weights instead)")
		if flagPartnerWeights == "" {
			flagPartnerWeights = deprecatedFlagPartnerStakes
		} else {
			log.Fatal().Msg("cannot use both --partner-stakes and --partner-weights flags (use only --partner-weights)")
		}
	}

	log.Info().Msg("collecting partner network and staking keys")
	partnerNodes, err := common.ReadFullPartnerNodeInfos(log, flagPartnerWeights, flagPartnerNodeInfoDir)
	if err != nil {
		log.Fatal().Err(err).Msg("failed to read full partner node infos")
	}
	log.Info().Msg("")

	log.Info().Msg("generating internal private networking and staking keys")
	internalNodes, err := common.ReadFullInternalNodeInfos(log, flagInternalNodePrivInfoDir, flagConfig)
	if err != nil {
		log.Fatal().Err(err).Msg("failed to read full internal node infos")
	}

	log.Info().Msg("")

	log.Info().Msg("checking constraints on consensus nodes")
	checkConstraints(partnerNodes, internalNodes)
	log.Info().Msg("")

	log.Info().Msg("assembling network and staking keys")
	stakingNodes := mergeNodeInfos(internalNodes, partnerNodes)
	log.Info().Msg("")

	// create flow.IdentityList representation of participant set
	participants := model.ToIdentityList(stakingNodes).Sort(flow.Canonical[flow.Identity])

	log.Info().Msg("reading root block data")
	block := readRootBlock()
	log.Info().Msg("")

	log.Info().Msg("reading root block votes")
	votes := readRootBlockVotes()
	log.Info().Msg("")

	log.Info().Msgf("received votes total: %v", len(votes))

	log.Info().Msg("reading dkg data")
	dkgData := readDKGData()
	log.Info().Msg("")

	log.Info().Msg("reading intermediary bootstrapping data")
	intermediaryData := readIntermediaryBootstrappingData()

	log.Info().Msg("constructing root QC")
	rootQC := constructRootQC(
		block,
		votes,
		model.FilterByRole(stakingNodes, flow.RoleConsensus),
		model.FilterByRole(internalNodes, flow.RoleConsensus),
		dkgData,
	)
	log.Info().Msg("")

	// if no root commit is specified, bootstrap an empty execution state
	if flagRootCommit == "0000000000000000000000000000000000000000000000000000000000000000" {
		commit := generateEmptyExecutionState(
			block.Header,
			intermediaryData.ExecutionStateConfig,
			participants,
		)
		flagRootCommit = hex.EncodeToString(commit[:])
	}

	log.Info().Msg("constructing root execution result and block seal")
	result, seal := constructRootResultAndSeal(flagRootCommit, block, intermediaryData.RootEpochSetup, intermediaryData.RootEpochCommit)
	log.Info().Msg("")

	// construct serializable root protocol snapshot
	log.Info().Msg("constructing root protocol snapshot")
	snapshot, err := inmem.SnapshotFromBootstrapStateWithParams(
		block,
		result,
		seal,
		rootQC,
		intermediaryData.ProtocolVersion,
		func(epochStateID flow.Identifier) (protocol_state.KVStoreAPI, error) {
			return kvstore.NewDefaultKVStore(
				intermediaryData.EpochCommitSafetyThreshold,
				intermediaryData.EpochExtensionViewCount,
				epochStateID)
		},
	)
	if err != nil {
		log.Fatal().Err(err).Msg("unable to generate root protocol snapshot")
	}

	// validate the generated root snapshot is valid
	verifyResultID := true
	err = badger.IsValidRootSnapshot(snapshot, verifyResultID)
	if err != nil {
		log.Fatal().Err(err).Msg("the generated root snapshot is invalid")
	}

	// validate the generated root snapshot QCs
	err = badger.IsValidRootSnapshotQCs(snapshot)
	if err != nil {
		log.Fatal().Err(err).Msg("root snapshot contains invalid QCs")
	}

	// write snapshot to disk
	err = common.WriteJSON(model.PathRootProtocolStateSnapshot, flagOutdir, snapshot.Encodable())
	if err != nil {
		log.Fatal().Err(err).Msg("failed to write json")
	}
	log.Info().Msgf("wrote file %s/%s", flagOutdir, model.PathRootProtocolStateSnapshot)
	log.Info().Msg("")

	// read snapshot and verify consistency
	rootSnapshot, err := loadRootProtocolSnapshot(model.PathRootProtocolStateSnapshot)
	if err != nil {
		log.Fatal().Err(err).Msg("unable to load serialized root protocol")
	}

	savedResult, savedSeal, err := rootSnapshot.SealedResult()
	if err != nil {
		log.Fatal().Err(err).Msg("could not load sealed result")
	}

	if savedSeal.ID() != seal.ID() {
		log.Fatal().Msgf("inconsistent seralization of the root seal: %v != %v", savedSeal.ID(), seal.ID())
	}

	if savedResult.ID() != result.ID() {
		log.Fatal().Msgf("inconsistent seralization of the root result: %v != %v", savedResult.ID(), result.ID())
	}

	if savedSeal.ResultID != savedResult.ID() {
		log.Fatal().Msgf("mismatch saved seal's resultID  %v and result %v", savedSeal.ResultID, savedResult.ID())
	}

	log.Info().Msg("saved result and seal are matching")

	err = badger.IsValidRootSnapshot(rootSnapshot, verifyResultID)
	if err != nil {
		log.Fatal().Err(err).Msg("saved snapshot is invalid")
	}

	// validate the generated root snapshot QCs
	err = badger.IsValidRootSnapshotQCs(snapshot)
	if err != nil {
		log.Fatal().Err(err).Msg("root snapshot contains invalid QCs")
	}

	log.Info().Msgf("saved root snapshot is valid")

	// copy files only if the directories differ
	log.Info().Str("private_dir", flagInternalNodePrivInfoDir).Str("output_dir", flagOutdir).Msg("attempting to copy private key files")
	if flagInternalNodePrivInfoDir != flagOutdir {
		log.Info().Msg("copying internal private keys to output folder")
		err := io.CopyDirectory(flagInternalNodePrivInfoDir, filepath.Join(flagOutdir, model.DirPrivateRoot))
		if err != nil {
			log.Error().Err(err).Msg("could not copy private key files")
		}
	} else {
		log.Info().Msg("skipping copy of private keys to output dir")
	}
	log.Info().Msg("")

	// print count of all nodes
	roleCounts := common.NodeCountByRole(stakingNodes)
	log.Info().Msg(fmt.Sprintf("created keys for %d %s nodes", roleCounts[flow.RoleConsensus], flow.RoleConsensus.String()))
	log.Info().Msg(fmt.Sprintf("created keys for %d %s nodes", roleCounts[flow.RoleCollection], flow.RoleCollection.String()))
	log.Info().Msg(fmt.Sprintf("created keys for %d %s nodes", roleCounts[flow.RoleVerification], flow.RoleVerification.String()))
	log.Info().Msg(fmt.Sprintf("created keys for %d %s nodes", roleCounts[flow.RoleExecution], flow.RoleExecution.String()))
	log.Info().Msg(fmt.Sprintf("created keys for %d %s nodes", roleCounts[flow.RoleAccess], flow.RoleAccess.String()))

	log.Info().Msg("🌊 🏄 🤙 Done – ready to flow!")
}

// readRootBlockVotes reads votes for root block
func readRootBlockVotes() []*hotstuff.Vote {
	var votes []*hotstuff.Vote
	files, err := common.FilesInDir(flagRootBlockVotesDir)
	if err != nil {
		log.Fatal().Err(err).Msg("could not read root block votes")
	}
	for _, f := range files {
		// skip files that do not include node-infos
		if !strings.Contains(f, model.FilenameRootBlockVotePrefix) {
			continue
		}

		// read file and append to partners
		var vote hotstuff.Vote
		err = common.ReadJSON(f, &vote)
		if err != nil {
			log.Fatal().Err(err).Msg("failed to read json")
		}

		votes = append(votes, &vote)
		log.Info().Msgf("read vote %v for block %v from signerID %v", vote.ID(), vote.BlockID, vote.SignerID)
	}
	return votes
}

// mergeNodeInfos merges the internal and partner nodes and checks if there are no
// duplicate addresses or node Ids.
//
// IMPORTANT: node infos are returned in the canonical ordering, meaning this
// is safe to use as the input to the DKG and protocol state.
func mergeNodeInfos(internalNodes, partnerNodes []model.NodeInfo) []model.NodeInfo {
	nodes := append(internalNodes, partnerNodes...)

	// test for duplicate Addresses
	addressLookup := make(map[string]struct{})
	for _, node := range nodes {
		if _, ok := addressLookup[node.Address]; ok {
			log.Fatal().Str("address", node.Address).Msg("duplicate node address")
		}
	}

	// test for duplicate node IDs
	idLookup := make(map[flow.Identifier]struct{})
	for _, node := range nodes {
		if _, ok := idLookup[node.NodeID]; ok {
			log.Fatal().Str("NodeID", node.NodeID.String()).Msg("duplicate node ID")
		}
	}

	// sort nodes using the canonical ordering
	nodes = model.Sort(nodes, flow.Canonical[flow.Identity])

	return nodes
}

// readRootBlock reads root block data from disc, this file needs to be prepared with
// rootblock command
func readRootBlock() *flow.Block {
	rootBlock, err := utils.ReadData[flow.Block](flagRootBlockPath)
	if err != nil {
		log.Fatal().Err(err).Msg("could not read root block data")
	}
	return rootBlock
}

// readDKGData reads DKG data from disc, this file needs to be prepared with
// rootblock command
func readDKGData() dkg.DKGData {
	encodableDKG, err := utils.ReadData[inmem.EncodableFullDKG](flagDKGDataPath)
	if err != nil {
		log.Fatal().Err(err).Msg("could not read DKG data")
	}

	dkgData := dkg.DKGData{
		PrivKeyShares: nil,
		PubGroupKey:   encodableDKG.GroupKey.PublicKey,
		PubKeyShares:  nil,
	}

	for _, pubKey := range encodableDKG.PubKeyShares {
		dkgData.PubKeyShares = append(dkgData.PubKeyShares, pubKey.PublicKey)
	}

	for _, privKey := range encodableDKG.PrivKeyShares {
		dkgData.PrivKeyShares = append(dkgData.PrivKeyShares, privKey.PrivateKey)
	}

	return dkgData
}

// Validation utility methods ------------------------------------------------

// loadRootProtocolSnapshot loads the root protocol snapshot from disk
func loadRootProtocolSnapshot(path string) (*inmem.Snapshot, error) {
	data, err := io.ReadFile(filepath.Join(flagOutdir, path))
	if err != nil {
		return nil, err
	}

	var snapshot inmem.EncodableSnapshot
	err = json.Unmarshal(data, &snapshot)
	if err != nil {
		return nil, err
	}

	return inmem.SnapshotFromEncodable(snapshot), nil
}

// readIntermediaryBootstrappingData reads intermediary bootstrapping data file from disk.
// This file needs to be prepared with rootblock command
func readIntermediaryBootstrappingData() *IntermediaryBootstrappingData {
	intermediaryData, err := utils.ReadData[IntermediaryBootstrappingData](flagIntermediaryBootstrappingDataPath)
	if err != nil {
		log.Fatal().Err(err).Msg("could not read root epoch data")
	}
	return intermediaryData
}

// generateEmptyExecutionState generates a new empty execution state with the
// given configuration. Sets the flagRootCommit variable for future reads.
func generateEmptyExecutionState(
	rootBlock *flow.Header,
	epochConfig epochs.EpochConfig,
	identities flow.IdentityList,
) (commit flow.StateCommitment) {

	log.Info().Msg("generating empty execution state")
	var serviceAccountPublicKey flow.AccountPublicKey
	err := serviceAccountPublicKey.UnmarshalJSON([]byte(flagServiceAccountPublicKeyJSON))
	if err != nil {
		log.Fatal().Err(err).Msg("unable to parse the service account public key json")
	}

	cdcInitialTokenSupply, err := cadence.NewUFix64(flagGenesisTokenSupply)
	if err != nil {
		log.Fatal().Err(err).Msg("invalid genesis token supply")
	}

	commit, err = run.GenerateExecutionState(
		filepath.Join(flagOutdir, model.DirnameExecutionState),
		serviceAccountPublicKey,
		rootBlock.ChainID.Chain(),
		fvm.WithRootBlock(rootBlock),
		fvm.WithInitialTokenSupply(cdcInitialTokenSupply),
		fvm.WithMinimumStorageReservation(fvm.DefaultMinimumStorageReservation),
		fvm.WithAccountCreationFee(fvm.DefaultAccountCreationFee),
		fvm.WithStorageMBPerFLOW(fvm.DefaultStorageMBPerFLOW),
		fvm.WithEpochConfig(epochConfig),
		fvm.WithIdentities(identities),
	)
	if err != nil {
		log.Fatal().Err(err).Msg("unable to generate execution state")
	}
	log.Info().Msg("")
	return commit
}

// validateOrPopulateEpochTimingConfig validates the epoch timing config flags. In case the
// `flagUseDefaultEpochTargetEndTime` value has been set, the function derives the values for
// `flagEpochTimingRefCounter`, `flagEpochTimingDuration`, and `flagEpochTimingRefTimestamp`
// from the configuration. Otherwise, it enforces that compatible values for the respective parameters have been
// specified (and errors otherwise). Therefore, after `validateOrPopulateEpochTimingConfig` ran,
// the targeted end time for the epoch can be computed via `rootEpochTargetEndTime()`.
// You can either let the tool choose default values, or specify a value for each config.
func validateOrPopulateEpochTimingConfig() error {
	// Default timing is intended for Benchnet, Localnet, etc.
	// Manually specified timings for Mainnet, Testnet, Canary.
	if flagUseDefaultEpochTargetEndTime {
		// No other flags may be set
		if !(flagEpochTimingRefTimestamp == 0 && flagEpochTimingDuration == 0 && flagEpochTimingRefCounter == 0) {
			return fmt.Errorf("invalid epoch timing config: cannot specify ANY of --epoch-timing-ref-counter, --epoch-timing-ref-timestamp, or --epoch-timing-duration if using default timing config")
		}
		flagEpochTimingRefCounter = flagEpochCounter
		flagEpochTimingDuration = flagNumViewsInEpoch
		flagEpochTimingRefTimestamp = uint64(time.Now().Unix()) + flagNumViewsInEpoch

		// compute target end time for initial (root) epoch from flags: `TargetEndTime = RefTimestamp + (RootEpochCounter - RefEpochCounter) * Duration`
		rootEpochTargetEndTimeUNIX := rootEpochTargetEndTime()
		rootEpochTargetEndTime := time.Unix(int64(rootEpochTargetEndTimeUNIX), 0)
		log.Info().Msgf("using default epoch timing config with root epoch target end time %s, which is in %s", rootEpochTargetEndTime, time.Until(rootEpochTargetEndTime))
	} else {
		// All other flags must be set
		// NOTE: it is valid for flagEpochTimingRefCounter to be set to 0
		if flagEpochTimingRefTimestamp == 0 || flagEpochTimingDuration == 0 {
			return fmt.Errorf("invalid epoch timing config: must specify ALL of --epoch-timing-ref-counter, --epoch-timing-ref-timestamp, and --epoch-timing-duration")
		}
		if flagEpochCounter < flagEpochTimingRefCounter {
			return fmt.Errorf("invalid epoch timing config: reference epoch counter must be less than or equal to root epoch counter")
		}

		// compute target end time for initial (root) epoch from flags: `TargetEndTime = RefTimestamp + (RootEpochCounter - RefEpochCounter) * Duration`
		rootEpochTargetEndTimeUNIX := rootEpochTargetEndTime()
		rootEpochTargetEndTime := time.Unix(int64(rootEpochTargetEndTimeUNIX), 0)
		log.Info().Msgf("using user-specified epoch timing config with root epoch target end time %s, which is in %s", rootEpochTargetEndTime, time.Until(rootEpochTargetEndTime))
	}
	return nil
}
